'''
Created on Oct 21, 2009

@author: mkiyer
'''
import os
import sys
import collections
import glob
import logging
from optparse import OptionParser
from bx.intervals.intersection import Interval, IntervalTree
from veggie.io.bed import bed_reader
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
import numpy as np
import operator

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.mlab as mlab
import math
    
def histogram(data, xlabel='x', ylabel='y', outfile=None):
    data = np.array(data)
    print 'average', np.average(data)

    fig = plt.figure()
    ax = fig.add_subplot(111)
    nbins=100
    # the histogram of the data
    n, bins, patches = ax.hist(data, bins=nbins, facecolor='green', alpha=0.5, log=False, normed=0)
    plt.axvline(x=np.average(data), linewidth=4, color='r')
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    #n, bins, patches = ax.hist(x2, bins=bins, facecolor='red', alpha=0.25)
    #l, = plt.plot(bins, mlab.normpdf(bins, 0.0, 1.0), 'r--', label='fit', linewidth=3)
    #legend([l, patches[0]], ['fit', 'hist'])
    ax.grid(True)
    if outfile is None:
        plt.show()
    else:
        plt.savefig(outfile)
        
    plt.close()
    return

def plot_scatter(data, 
                 outfile=None,
                 xlabel='Control Group',
                 ylabel='Experimental Group',
                 title='Coverage scatter plot'):  
    fig = plt.figure()
    ax = fig.add_subplot(111)
    plt.title(title)

    #sorted_data = sorted(data, operator.itemgetter(0))
    
    xdata = [d[0] for d in data]
    ydata = [d[1] for d in data]
    # size?
    # s = 20*np.log(np.array([x.ratio for x in intervals]))        
    # color?
    # c = [x.category * cm.jet.N / len(category_names)] * len(intervals)
    # c = x.category * cm.jet.N / len(category_names)
    plt.scatter(xdata, ydata, 
                s=30, c='b', marker='o', cmap=None, norm=None,
                vmin=None, vmax=None, alpha=1.0, linewidths=None,
                verts=None)

    if outfile is None:
        plt.show()
    else:
        plt.savefig(outfile)
    plt.close()
    
    return    
    #ax.set_xscale('log')
    #ax.set_yscale('log')

    plt.axvline(x=np.average(xdata), linewidth=4, color='r')
    plt.axhline(y=np.average(ydata), linewidth=4, color='g')
    
    print 'x', np.average(xdata)
    print 'y', 'avg', np.average(ydata), 'min', np.min(ydata), 'max', np.max(ydata)
    
    # label=category_names[category])
    #plt.legend()
    # best fit line
    sorted_data = sorted(data, key=operator.itemgetter(0))
    x, y = [d[0] for d in sorted_data], [d[1] for d in sorted_data]    
    m = np.polyfit(x, y, 1)
    yfit = np.polyval(m, x)
    plt.plot(x, yfit, 'g')
    # correlation coefficient
    c = np.corrcoef(xdata, ydata)[0,1]
    r2 = c**2
    plt.figtext(0.815, 0.013, ' r^2=%.3f' % r2, color='black', weight='roman',
               size='small')    
    plt.grid()
    themax = max(x[-1], y[-1])
    plt.axis([-10, themax, -10, themax])
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)


def clustertree(intervals, cluster_distance, min_size=1):  
    from bx.intervals.cluster import ClusterTree
    # arguments to ClusterTree are:  
    # - Distance in basepairs for two reads to be in the same cluster;  
    #   for instance 20 would group all reads with 20bp of each other  
    # - Number of reads necessary for a group to be considered a cluster;  
    #   2 returns all groups with 2 or more overlapping reads  
    cluster_trees = collections.defaultdict(lambda: ClusterTree(cluster_distance, min_size))    
    id = 0
    for chrom, start, end in intervals:
        cluster_trees[chrom].insert(start, end, id)
        id += 1
    result = []
    for chrom, cluster_tree in cluster_trees.items():
        for start, end, ids in cluster_tree.getregions():
            result.append((chrom, start, end))     
    return result


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    outfhd = sys.stdout
    percent_overlap = 50.0
    results_suffix = "unknown_candidates.bed"
    qc_fail_libraries = set(['30WU2AAXX_2',
                             '30WU2AAXX_3',
                             '42P6UAAXX_8'])

    # parse command line
    optionparser = OptionParser("usage: %prog [options] <sample_xml_file> <bed_dir>")
    (options, args) = optionparser.parse_args()    
    sample_groups = list(parse_samplegroups_xml(args[0]))
    samples = set([])
    for sgroup in sample_groups:
        samples.update(sgroup.samples)
    results_path = args[1]


    trees = collections.defaultdict(lambda: IntervalTree())
    annots = []
    libraries = []
        
    for sample in samples:
        for res in get_sampledb().params[sample]['lanes']:
            flowcell, lane = res
            library_id = '%s_%d' % (flowcell, lane)
            library_bedfile = '%s.%s' % (library_id, results_suffix)
            library_path = os.path.join(results_path, library_bedfile)
            if not os.path.exists(library_path):
                logging.critical('File not found: %s' % library_path)
            if library_id in qc_fail_libraries:
                logging.debug('Skipping poor quality library %s' % (library_id))
                continue
            logging.debug('Adding library %s to sample %s' % (library_id, sample))
            libraries.append(library_id)
            for annot in bed_reader(open(library_path)):
                # add sample attribute
                trees[annot.chrom].insert_interval(Interval(annot.start, annot.end, value=annot))
                annots.append(annot)

    intervals = [(x.chrom, x.start, x.end) for x in annots]
    cluster_results = clustertree(intervals, 0, 6)

    print len(cluster_results)

    sys.exit(0)
    
    avg_percent_overlaps = []
    overlap_vs_cov = []
    recurrences = []
    # intersect each interval
    for a1 in annots:
        id1 = a1.name.split('.')[0] 
        hits = trees[a1.chrom].find(a1.start, a1.end)
        true_hits = []
        library_hits = collections.defaultdict(lambda: [])
                    
        for hit in hits:
            a2 = hit.value
            id2 = a2.name.split('.')[0]                
            # ignore hits from same library (isoforms)
            if id1 == id2:
                continue
            # find overlapping interval
            ostart, oend = sorted((a1.start, a1.end, a2.start, a2.end))[1:3]
            percent = max(float(oend - ostart)/(a1.end - a1.start),
                          float(oend - ostart)/(a2.end - a2.start))
            # TODO: set an overlap threshold?
            library_hits[id2].append((a2,percent))
        
        # choose the isoform with best overlap for each library
        best_percents = []
        best_covs = [a1.score]
        for l, isoform_hits in library_hits.iteritems():                
            best_hit = sorted(isoform_hits, 
                              key=operator.itemgetter(1), 
                              reverse=True)[0]
            best_covs.append(best_hit[0].score)
            best_percents.append(best_hit[1])
        recurrences.append(len(library_hits))
        avg_percent_overlap = np.sum(best_percents)/float(len(libraries)-1)
        avg_cov = np.sum(best_covs)/float(len(libraries))            
        overlap_vs_cov.append((avg_percent_overlap, avg_cov))
        avg_percent_overlaps.append(avg_percent_overlap)

    histogram(recurrences)
    #histogram(avg_percent_overlaps)
    plot_scatter(overlap_vs_cov, 
                 outfile='overlap_vs_cov.png',
                 xlabel='% Overlap',
                 ylabel='Average Coverage',
                 title='%d' % (len(libraries))) 
        
    sys.exit(0)
    
    for sample in samples:

        trees = collections.defaultdict(lambda: IntervalTree())
        annots = []
        libraries = []
                
        for res in get_sampledb().params[sample]['lanes']:
            flowcell, lane = res
            library_id = '%s_%d' % (flowcell, lane)
            library_bedfile = '%s.%s' % (library_id, results_suffix)
            library_path = os.path.join(results_path, library_bedfile)
            if not os.path.exists(library_path):
                logging.critical('File not found: %s' % library_path)
            if library_id in qc_fail_libraries:
                logging.debug('Skipping poor quality library %s' % (library_id))
                continue
            logging.debug('Adding library %s to sample %s' % (library_id, sample))
            libraries.append(library_id)
            for annot in bed_reader(open(library_path)):
                # add sample attribute
                trees[annot.chrom].insert_interval(Interval(annot.start, annot.end, value=annot))
                annots.append(annot)

        avg_percent_overlaps = []
        overlap_vs_cov = []
        # intersect each interval
        for a1 in annots:
            id1 = a1.name.split('.')[0] 
            hits = trees[a1.chrom].find(a1.start, a1.end)
            true_hits = []
            library_hits = collections.defaultdict(lambda: [])
                        
            for hit in hits:
                a2 = hit.value
                id2 = a2.name.split('.')[0]                
                # ignore hits from same library (isoforms)
                if id1 == id2:
                    continue
                # find overlapping interval
                ostart, oend = sorted((a1.start, a1.end, a2.start, a2.end))[1:3]
                percent = max(float(oend - ostart)/(a1.end - a1.start),
                              float(oend - ostart)/(a2.end - a2.start))
                library_hits[id2].append((a2,percent))
            
            # choose the isoform with best overlap for each library
            best_percents = []
            best_covs = [a1.score]
            for l, isoform_hits in library_hits.iteritems():                
                best_hit = sorted(isoform_hits, 
                                  key=operator.itemgetter(1), 
                                  reverse=True)[0]
                best_covs.append(best_hit[0].score)
                best_percents.append(best_hit[1])
            avg_percent_overlap = np.sum(best_percents)/float(len(libraries)-1)
            avg_cov = np.sum(best_covs)/float(len(libraries))            
            overlap_vs_cov.append((avg_percent_overlap, avg_cov))
            avg_percent_overlaps.append(avg_percent_overlap)

        #histogram(avg_percent_overlaps)
        plot_scatter(overlap_vs_cov, 
                     outfile='%s_overlap_vs_cov.png' % sample,
                     xlabel='% Overlap',
                     ylabel='Average Coverage',
                     title='%s %d' % (sample, len(libraries))) 
